import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt

# 生成正弦波数据
def generate_sine_wave(n_steps):
    time = np.linspace(0, 2 * np.pi, n_steps)
    sine_wave = np.sin(time)
    return sine_wave

# 准备数据
def prepare_data(sine_wave, seq_length):
    inputs = []
    targets = []
    for i in range(len(sine_wave) - seq_length):
        inputs.append(sine_wave[i:i+seq_length])
        targets.append(sine_wave[i+seq_length])
    inputs = np.array(inputs)
    targets = np.array(targets)
    inputs = torch.from_numpy(inputs).float().unsqueeze(2)
    targets = torch.from_numpy(targets).float().unsqueeze(1)
    return inputs, targets

# 定义 RNN 模型
class SimpleRNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleRNN, self).__init__()
        self.rnn = nn.RNN(input_size, hidden_size, batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        out, _ = self.rnn(x) # (batch,seq_length,hidden_size)
        out = self.fc(out[:, -1, :]) # (batch,1,hidden_size) # 去最后一个输出结果作为最终的结果
        return out  

# 参数设置
n_steps = 1000
seq_length = 20
input_size = 1
hidden_size = 32
output_size = 1
learning_rate = 0.01
num_epochs = 100

# 生成数据
sine_wave = generate_sine_wave(n_steps)
inputs, targets = prepare_data(sine_wave, seq_length)  #  inputs torch.Size([980, 20, 1]) targets torch.Size([980, 1])

# 初始化模型、损失函数和优化器
model = SimpleRNN(input_size, hidden_size, output_size)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

# 训练模型
for epoch in range(num_epochs):
    outputs = model(inputs) # outputs (980,1)
    loss = criterion(outputs, targets)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if (epoch + 1) % 10 == 0:
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

# 进行预测
test_input = inputs[-1].unsqueeze(0) #(1,20,1)
predictions = []
for _ in range(100):
    prediction = model(test_input) # (1,1)
    predictions.append(prediction.item())
    test_input = torch.cat((test_input[:, 1:, :], prediction.unsqueeze(2)), dim=1) # (1,1,1)

# 绘制结果
plt.plot(sine_wave, label='Original Sine Wave')
plt.plot(np.arange(len(sine_wave), len(sine_wave) + len(predictions)), predictions, label='Predicted Sine Wave')
plt.legend()
plt.show()
    